import numpy as np
import operator


def quick_sort(lists, left, right):
    # 快速排序
    if left >= right:
        return lists
    key = lists[left]
    low = left
    high = right
    while left < right:
        while left < right and lists[right][1] >= key[1]:
            right -= 1
        lists[left] = lists[right]
        while left < right and lists[left][1] <= key[1]:
            left += 1
        lists[right] = lists[left]
    lists[right] = key
    quick_sort(lists, low, left - 1)
    quick_sort(lists, left + 1, high)
    return lists


def python_argsort(array):
    length = array.shape[0]
    lists = [(i, array[i]) for i in range(length)]
    lists = quick_sort(lists, 0, length-1)
    index_sorted = [x[0] for x in lists]
    return index_sorted


def classify0(inX: list, dataSet: np.ndarray, labels: list, k: int) -> object:
    def dist_generator(vec1, vec2):
        dist = np.sqrt(np.sum(np.square(vec1 - vec2)))
        return dist

    dataSet = dataSet.astype(np.float32)
    dataSetSize = dataSet.shape[0]

    aim_data = np.array(inX).astype(np.float32)
    distances = np.zeros((dataSetSize,))
    for i in range(dataSetSize):
        distances[i] = dist_generator(aim_data, dataSet[i])

    sortedDistances = python_argsort(distances)

    classCount = {}
    for i in range(k):
        numOflabel = labels[sortedDistances[i]]
        classCount[numOflabel] = classCount.get(numOflabel, 0) + 1
    sortedClassCount = sorted(classCount.items(), key=operator.itemgetter(1), reverse=True)
    return sortedClassCount[0][0]


def createDataSet():
    group = np.array([[1.0, 1.1], [1.0, 1.0], [0, 0], [0, 0.1]])
    labels = ['A', 'A', 'B', 'B']
    return group, labels


import matplotlib.pyplot as plt

if __name__ == "__main__":
    dataSet, labels = createDataSet()
    print(classify0([0, 0], dataSet, labels, 3))
    fig = plt.figure()
    ax = fig.add_subplot(111)
    ax.scatter(dataSet[:, 0], dataSet[:, 1])
    plt.show()
    var = dataSet.shape
    print(var)
